#### SIMULATE DLES AND TEST FOR INFLATION WITH
#### STEPHENSON TEST
#### This is tricky because including a positive prevalence rate 
#### Leads to a positive test statistic
#### Start by documenting the magnitude of false positive rate
#### Packages ####
library(tidyverse) # data manipulation
library(data.table) # high performance data.frame operations
library(DeclareDesign) # simulation pipeline
library(future) # parallel processing
library(here) # relative path management
library(coin) # conditional independence testing
# devtools::install_github("li-xinran/RIQITE")
library(RIQITE) ## A wrapper for Stephenson signed rank test in coin package

#### GLOBAL PARAMETERS ####
# Grid for subset sizes to compute in Stephenson test
subsets = c(2, 5, 10, 50, 100)

#### FUNCTION TO OUTPUT DECLARE-DESIGN FRIENDLY STEPHENSON TEST ####
# Inner function computes stephenson signed rank sum with p-value
# Tidy-friendly output
st = function(m, data){
  temp = stephenson_test(formula = Y_diff ~ Z_a, data = data, subset_size = m,
                         alternative = "greater")
  
  out = data.frame(statistic = temp@statistic@teststatistic,
                   p.value = pvalue(temp)[1], m = m)
  
  return(out)
}

# Outer function repeats the test for an arbitrary number of subset sizes
# Output is a data frame of statistic, p.value, and m
st_test = function(data){
  temp = lapply(subsets, st, data = data)
  
  out = data.frame(rbindlist(temp))
  
  return(out)
}

## Test for equality of prevalence estimates
## This is the same  as the difference in differences
## This needs the transformed version of the data
## So place after that step in the designer
dd = function(data){
  temp = tidy(lm_robust(Y ~ Z*list, data  = data, clusters = ID, se_type = "stata")) %>% 
    filter(term == "Z:listb")
  
  out = data.frame(statistic = temp$estimate,
                   p.value = temp$p.value)
  
  return(out)
}

#### EXPANDABLE DESIGN ####
#### Parameters ####
## N = sample size
## delta = inflation/deflation rate. May need  to rethink this
## rho = correlation between baseline lists
designer = function(N, gamma, rho){
  # Model
  # The key is to think of altered responses as an alternative DGP, 
  # then reveal outcomes accordingly
  model = declare_model(
    N = N,
    count_a = rbinom(N, size = 4, prob = 0.5),
    count_b = correlate(
      given = count_a,
      rho = rho, # correlation between baseline lists
      rbinom,
      size = 4,
      prob = 0.5),
    # make magnitude of change completely unrelated to mimic having no reason why respondents engage with lists differently
    delta_a = sample(1:2, size = N, replace= TRUE),
    delta_b= sample(1:2, size = N, replace= TRUE),
    # X is whether they have sensitive trait
    # Based on Blair et al 2020
    # A single LE would not detect this effect size
    # but a DLE would have enough power at 80%
    # Change this for appendix simulations
    X = rbinom(N, size = 1, prob = 0.15),
    # Whether they deflate response when they see sensitive item
    inflate = rbinom(N, size = 1, prob = gamma),
    potential_outcomes(Y_a0 ~ X * Z_a + count_a,
                       conditions = list(Z_a = c(0,1))),
    potential_outcomes(Y_b0 ~ X * Z_b + count_b,
                       conditions = list(Z_b = c(0,1))),
    potential_outcomes(Y_a1 ~ (X * Z_a + count_a) + delta_a*inflate,
                       conditions = list(Z_a = c(0,1))),
    potential_outcomes(Y_b1 ~ (X * Z_b + count_b) + delta_b*inflate,
                       conditions = list(Z_b = c(0,1)))
  )
  
  # Inquiry
  inquiry = declare_inquiry(prevalence_rate = mean(X),
                            cor_ab = cor(count_a, count_b)) # I think this is a placeholder
  
  # Data Strategy
  assign = declare_assignment(Z_a = complete_ra(N),
                              Z_b = 1 - Z_a) 
  
  measurement = declare_measurement(Y_a0 = reveal_outcomes(Y_a0 ~ Z_a),
                                    Y_b0 = reveal_outcomes(Y_b0 ~ Z_b),
                                    Y_a1 = reveal_outcomes(Y_a1 ~ Z_a),
                                    Y_b1 = reveal_outcomes(Y_b1 ~ Z_b),
                                    # Responses can't be lower than 0
                                    Y_a = ifelse(Z_a == 1, 
                                                 pmin(5, Y_a1),
                                                 pmin(4, Y_a0)),
                                    Y_b = ifelse(Z_a == 1, 
                                                 pmin(4, Y_b1), 
                                                 pmin(5, Y_b1)),
                                    # compute diff in paired responses
                                    Y_diff = Y_a - Y_b)
  
  # Answer strategy
  stephenson = declare_test(handler = label_test(st_test), 
                            label = "stephenson")
  
  single_a = declare_estimator(Y_a ~ Z_a, inquiry = "prevalence_rate", label = "List A")
  
  single_b = declare_estimator(Y_b ~ Z_b, inquiry = "prevalence_rate", label = "List B")
  
  dle_handler = declare_step(handler = function(data){
    data %>% 
      pivot_longer(cols = c(Z_a, Z_b, Y_a, Y_b)) %>% 
      separate(name, into = c("variable", "list")) %>% 
      pivot_wider(names_from = "variable", values_from = "value")}
  )
  
  prev_diff = declare_test(handler = label_test(dd), 
                           label = "diff_in_prevalence")
  
  dle = declare_estimator(Y ~ Z + list, clusters = ID, se_type = "stata", inquiry = "prevalence_rate", label = "DLE")
  
  model + inquiry + assign + measurement + stephenson +
    single_a + single_b + dle_handler + prev_diff + dle
  
}

designs = expand_design(designer,
                        N = 1000, # sample size
                        gamma = seq(0, 1, by = 0.1), # prop deflating
                        rho =  c(0, 0.4, 0.8), # correlation between baseline lists
                        expand = TRUE)



#### SIMULATE ####
plan(multicore) # ignored in Windows or RStudio
set.seed(20220126)
sims_inf = simulate_design(designs, sims = 1000)
plan(sequential)

#### SAVE ####
save(sims_inf, file = here("dle_sims_inflation.rda"))
